#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
代理服务，用于拦截和重放请求
"""

import logging
import socket
import struct
import threading
import time
from flask import Flask, request, Response
import requests
from intercept.services.intercept_service import InterceptService
from intercept.services.iptables_manager import IptablesManager

logger = logging.getLogger(__name__)

class ProxyService:
    """代理服务"""
    
    # 代理服务器
    _app = None
    
    # 代理服务器线程
    _server_thread = None
    
    # 是否正在运行
    _is_running = False
    
    # 代理端口
    _proxy_port = 5555
    
    @staticmethod
    def create_app():
        """
        创建 Flask 应用
        
        Returns:
            Flask: Flask 应用
        """
        app = Flask(__name__)
        
        @app.route('/', defaults={'path': ''}, methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS'])
        @app.route('/<path:path>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'HEAD', 'OPTIONS'])
        def proxy(path):
            """代理请求"""
            try:
                # 获取原始目标信息
                original_dst = ProxyService._get_original_destination(request.environ)
                if not original_dst:
                    return Response("无法获取原始目标信息", status=500)
                
                original_host, original_port = original_dst
                
                # 构建原始 URL
                scheme = "https" if original_port == 443 else "http"
                url = f"{scheme}://{original_host}"
                if (scheme == "http" and original_port != 80) or (scheme == "https" and original_port != 443):
                    url += f":{original_port}"
                
                if path:
                    url += f"/{path}"
                
                # 获取请求信息
                method = request.method
                headers = {k: v for k, v in request.headers.items() if k.lower() != 'host'}
                headers['Host'] = f"{original_host}:{original_port}"
                body = request.get_data()
                
                # 识别服务
                service_name = ProxyService._identify_service(original_host, original_port)
                if not service_name:
                    # 如果无法识别服务，直接转发请求
                    return ProxyService._forward_request(method, url, headers, body)
                
                # 拦截请求
                intercepted, status_code, response_headers, response_body, record_id = InterceptService.intercept_request(
                    service_name=service_name,
                    method=method,
                    url=url,
                    headers=headers,
                    body=body
                )
                
                # 构建响应
                if intercepted:
                    # 添加拦截标记
                    response_headers['X-Intercepted'] = 'true'
                    logger.info(f"已拦截请求: {method} {url}")
                else:
                    logger.info(f"已转发请求: {method} {url}")
                
                return Response(
                    response_body,
                    status=status_code,
                    headers=response_headers
                )
            
            except Exception as e:
                logger.error(f"代理请求失败: {e}")
                return Response(f"代理请求失败: {str(e)}", status=500)
        
        return app
    
    @staticmethod
    def start(port=5555):
        """
        启动代理服务
        
        Args:
            port: 代理端口
            
        Returns:
            bool: 是否成功启动
        """
        if ProxyService._is_running:
            logger.warning("代理服务已经在运行")
            return False
        
        # 保存代理端口
        ProxyService._proxy_port = port
        
        # 创建 Flask 应用
        ProxyService._app = ProxyService.create_app()
        
        # 启动代理服务器线程
        ProxyService._is_running = True
        ProxyService._server_thread = threading.Thread(
            target=ProxyService._app.run,
            kwargs={
                'host': '0.0.0.0',
                'port': port,
                'debug': False,
                'use_reloader': False,
                'threaded': True
            },
            daemon=True
        )
        ProxyService._server_thread.start()
        
        # 等待服务器启动
        time.sleep(1)
        
        logger.info(f"已启动代理服务，端口: {port}")
        return True
    
    @staticmethod
    def stop():
        """
        停止代理服务
        
        Returns:
            bool: 是否成功停止
        """
        if not ProxyService._is_running:
            logger.warning("代理服务未在运行")
            return False
        
        # 停止代理服务器
        ProxyService._is_running = False
        
        # 关闭所有请求
        if ProxyService._app:
            # 关闭 Flask 应用
            func = request.environ.get('werkzeug.server.shutdown')
            if func is not None:
                func()
        
        # 等待线程结束
        if ProxyService._server_thread:
            ProxyService._server_thread.join(timeout=5)
            ProxyService._server_thread = None
        
        logger.info("已停止代理服务")
        return True
    
    @staticmethod
    def is_running():
        """
        检查是否正在运行
        
        Returns:
            bool: 是否正在运行
        """
        return ProxyService._is_running
    
    @staticmethod
    def get_port():
        """
        获取代理端口
        
        Returns:
            int: 代理端口
        """
        return ProxyService._proxy_port
    
    @staticmethod
    def _get_original_destination(environ):
        """
        获取原始目标信息
        
        Args:
            environ: WSGI 环境
            
        Returns:
            tuple: (host, port) 原始目标主机和端口
        """
        try:
            # 获取客户端连接的文件描述符
            if 'wsgi.input' not in environ:
                return None
            
            # 获取 socket 对象
            sock = socket.fromfd(
                environ['wsgi.input'].fileno(),
                socket.AF_INET,
                socket.SOCK_STREAM
            )
            
            # 获取原始目标信息
            SO_ORIGINAL_DST = 80  # 这是 Linux 特定的 socket 选项
            dst = sock.getsockopt(socket.SOL_IP, SO_ORIGINAL_DST, 16)
            port, raw_ip = struct.unpack_from('!2xH4s', dst)
            ip = socket.inet_ntoa(raw_ip)
            
            return (ip, port)
        
        except Exception as e:
            logger.error(f"获取原始目标信息失败: {e}")
            return None
    
    @staticmethod
    def _identify_service(host, port):
        """
        识别服务
        
        Args:
            host: 主机
            port: 端口
            
        Returns:
            str: 服务名称，如果无法识别则返回 None
        """
        from db.mongo import get_db
        
        try:
            # 查询匹配的服务
            db = get_db()
            service = db.intercept_services.find_one({
                "service_ip": host,
                "service_port": port,
                "is_active": True
            })
            
            if service:
                return service["service_name"]
            
            # 如果没有精确匹配，尝试只匹配 IP
            service = db.intercept_services.find_one({
                "service_ip": host,
                "is_active": True
            })
            
            if service:
                return service["service_name"]
            
            return None
        
        except Exception as e:
            logger.error(f"识别服务失败: {e}")
            return None
    
    @staticmethod
    def _forward_request(method, url, headers, body):
        """
        转发请求
        
        Args:
            method: 请求方法
            url: 请求URL
            headers: 请求头
            body: 请求体
            
        Returns:
            Response: 响应对象
        """
        try:
            # 发送请求
            response = requests.request(
                method=method,
                url=url,
                headers=headers,
                data=body,
                allow_redirects=False,
                timeout=30
            )
            
            # 返回响应
            return Response(
                response.content,
                status=response.status_code,
                headers=dict(response.headers)
            )
        
        except Exception as e:
            logger.error(f"转发请求失败: {e}")
            return Response(f"转发请求失败: {str(e)}", status=500)
